#!/usr/bin/python
# coding=utf-8

import traceback
import psycopg2
import time
import sys
import sqlparse
from sqlparse.sql import Identifier, Where, Comparison, Parenthesis, Function

#处理UnicodeDecodeError: 'ascii' codec can't decode byte 0xe9 in position 0: ordinal not in range(128)
default_encoding = 'utf-8'
if sys.getdefaultencoding() != default_encoding:
    reload(sys)
    sys.setdefaultencoding(default_encoding)

# 连接即将执行query的数据库
def conn_db(order_id):
    # print '[{0}] 连接数据库 ...'.format(time.strftime('%X'))
    #连接部署数据库，获取目标执行数据库的连接方式
    global conn
    conn = psycopg2.connect("host=opsdb.yazuoyw.com dbname=ops port=5432 user=opser password=opsEr3469")
    global cur
    cur = conn.cursor()
    return conn

# 以元组的形式获取上传的query语句，并过滤掉换行符
def get_query(order_id):
    conn = conn_db(order_id)
    query_sql = 'select script from sql_contents where sql_order_id=%s;'
    cur.execute(query_sql, (order_id,))
    conn.commit()
    query_lists, = cur.fetchall()
    query_list = sqlparse.parse(query_lists[0].strip('\n').strip('\r\n').strip('\r'))
    return query_list

# 对query语句进行分类统计，并返回校验结果
def analyze_query(order_id):
    f = open('./python_report.log','a')
    f.write('[{0}] 开始分析上线单--%s ...\n'.format(time.strftime('%X')) %order_id)
    f.write('[{0}] 读取SQL语句...\n'.format(time.strftime('%X')))
    query_list = get_query(order_id)
    sum_u = 0
    sum_i = 0
    sum_d = 0
    sum_s = 0
    sum_o = 0
    sum_p = 0
    sum_a = len(query_list)
    report_list = ''
    report = ''
    f.write('[{0}] 正在分析SQL...\n'.format(time.strftime('%X')))
    for query in query_list:
        sum_p = sum_p + 1
        #判断query类别，不同类别的query调用不同的方法来完成校验
        query_dml = query.get_type()
        if query_dml == "INSERT":
            sum_i = sum_i + 1
            report = insert_query(query)
        elif query_dml == "UPDATE":
            sum_u = sum_u + 1
            report = update_query(order_id, query)
        elif query_dml == "DELETE":
            sum_d = sum_d + 1
            report = delete_query(order_id, query)
        elif query_dml == "UNKNOWN":
            continue
#        elif query_dml == "SELECT":
#            sum_s = sum_s + 1
        else:
            sum_o += 1
            report = "发现非insert,update,delete SQL:%s" %query

        report_list = report_list + '\n' + report

        parse_done = "update sql_orders set done_entries=%s where id=%s;"
        cur.execute(parse_done,(sum_p,order_id))
        conn.commit()

    #调用返回校验报告的方法
    f.write('[{0}] 返回校验报告...\n'.format(time.strftime('%X')))
    parse_report(sum_a,sum_i, sum_u, sum_d, sum_s, sum_o, report_list,order_id)
    conn.close()
    f.closed

# INSERT语句处理规则：只判断模式名和关键字是否而存在
def insert_query(query):
    into_key_word = ''
    values_key_work = ''
    tablename = ''
    for token in query.tokens:
        #查找表名并且输出
        if isinstance(token, Identifier) and tablename == '':
            tablename = str(token).strip(' ').split('(')[0]

        #判断关键字如果存在就输出，否则为空
        if token.is_keyword:
            if str(token) == 'into' or str(token) == 'INTO':
                into_key_word = token
            if str(token) == 'values' or str(token) == 'VALUES':
                values_key_work = token

    #表名和关键字有任何一个不存在，就认为sql校验不通过
    if (tablename is '') or (into_key_word is '') or (values_key_work is ''):
        report = '模式名或关键字不存在:%s' %query
    else:
        report = ''
    return report

# UPDATE语句处理规则：需要判断update query关键字是否存在，模式名是否存在，表和主键的对应关系是否正确
def update_query(order_id, query):
    key_word_lists = ['and','or','between','select','update','in','exists']
    primary_key = ''
    set_key_word = ''
    tablename = ''
    report_tmp1 = ''
    for token in query.tokens:
        if isinstance(token, Identifier) and tablename == '':
            tablename = str(token).strip(' ').split('(')[0]
        if token.is_keyword:
            if str(token) == 'set' or str(token) == 'SET':
                set_key_word = token
        if isinstance(token, Where):
            # primary_key = str(token).split('=')[0].split(' ')[-1]
            # print(primary_key)
            for key_all in token.tokens:
                if isinstance(key_all,Comparison) and primary_key is '':
                    primary_key = str(key_all).split('=')[0].strip(' ')
                if str(key_all) in key_word_lists:
                    report_tmp1 = '检索条件过多：%s' %query
    if (primary_key is not '') and (tablename is not '') and (set_key_word is not ''):
        report_tmp2 = primary_dict(primary_key,tablename,query,order_id)
    else:
        report_tmp2 = '模式名或关键字不存在:%s' %query

    report =  report_tmp1 + '\n' + report_tmp2
    return report

# DELETE语句处理规则：需要判断delete query关键字是否存在，模式名是否存在，表和主键的对应关系是否正确
def delete_query(order_id, query):
    key_word_lists = ['and','or','between','select','update','in','exists']
    primary_key = ''
    from_key_word = ''
    tablename = ''
    report_tmp1 = ''
    for token in query.tokens:
        if isinstance(token, Identifier) and tablename == '':
            tablename = str(token).strip(' ').split('(')[0]
        if token.is_keyword:
            if str(token) == 'from' or str(token) == 'FROM':
                from_key_word = token
        if isinstance(token, Where):
            # primary_key = str(token).split('=')[0].split(' ')[-1]
            for key_all in token.tokens:
                if isinstance(key_all,Comparison) and primary_key is '':
                    primary_key = str(key_all).split('=')[0].strip(' ')
                if str(key_all) in key_word_lists:
                    report_tmp1 = '检索条件过多：%s' %query
    if (primary_key is not '') and (tablename is not '') and (from_key_word is not ''):
        report_tmp2 = primary_dict(primary_key,tablename,query,order_id)
    else:
        report_tmp2 = '模式名或关键字不存在:%s' %query
    report =  report_tmp1 + '\n' + report_tmp2
    return report

#判断表与主键的关系是否匹配
def primary_dict(primary_key,tablename,query,order_id):
    sql = "select primary_key from primary_relation where tablename='%s'" %tablename

    #判断sql执行是否正常，如果正常需要对表名和主键进行匹配，否则直接退出
    try:
        cur.execute(sql)
        value, = cur.fetchall()
        primary_key_tmp = value[0]
        if str(primary_key_tmp) == str(primary_key):
            report = ''
        else:
            report = '表名与主键不匹配:%s' %query
    except:
        report = '表名与主键不匹配:%s' %query
    return report

#返回校验报告及上线单状态
def parse_report(sum_a,sum_i, sum_u, sum_d, sum_s, sum_o, report_list,order_id):
    #如果report_list不为空或sum_o不等于0，说明query有问题，需要输出详细报告并更改申请单状态为rejected
    if report_list.strip(' ').strip('\n') is not '' or sum_o != 0:
        checkreport = ' insert语句 :%s条 \n update语句 :%s条 \n delete语句 :%s条 \n 非DML语句   :%s条 \n report: %s' %(sum_i, sum_u, sum_d,  sum_o, report_list)
        report_sql = "update sql_contents set report=%s where sql_order_id=%s;"
        status_sql = "update sql_orders set done_entries=0,status='rejected',all_entries=%s where id=%s"
        cur.execute(report_sql, (checkreport, order_id))
        conn.commit()
        cur.execute(status_sql, (sum_a, order_id))
        conn.commit()
    #否则说明query没有问题，需要输出详细报告并更改申请单状态为ready
    else:
        checkreport = ' insert语句 :%s条 \n update语句 :%s条 \n delete语句 :%s条 \n 非DML语句  :%s条 \n report: success' %(sum_i, sum_u, sum_d, sum_o)
        report_sql = "update sql_contents set report=%s where sql_order_id=%s;"
        status_sql = "update sql_orders set done_entries=0,status='ready',all_entries=%s where id=%s"
        cur.execute(report_sql, (checkreport, order_id))
        conn.commit()
        cur.execute(status_sql, (sum_a, order_id))
        conn.commit()

# 抛出脚本执行异常信息，并返回到报告中
def except_fun(order_id, execpt):
    print('[{0}] 执行parse_query脚本报错...'.format(time.strftime('%X')))
    conn = conn_db(order_id)
    report_text = "%s" %execpt
    report_sql = "update sql_contents set report=%s where sql_order_id=%s;"
    cur.execute(report_sql, (report_text, order_id))
    conn.commit()
    conn.close()

if __name__ == "__main__":
    try:
        analyze_query(sys.argv[1])
    except Exception,e:
        exstr = traceback.format_exc()
        except_fun(sys.argv[1],exstr)

